import argparse
import time
import os

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from utils import plot_metrics
from dataset import get_cifar100_loaders  
from utils import profile_model  
from utils import find_lr 


from model.eca_resnet import eca_resnet18
from model.eca_resnet import eca_resnet50



def train(train_loader, model, criterion, optimizer, epoch, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for i, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)

        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * inputs.size(0)
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    avg_loss = running_loss / total
    acc = correct / total
    return avg_loss, acc


def test(test_loader, model, criterion, device):
    model.eval()
    total = 0
    correct = 0
    loss_total = 0.0

    with torch.no_grad():
        for inputs, targets in test_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)

            loss_total += loss.item() * inputs.size(0)
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

    avg_loss = loss_total / total
    acc = correct / total
    return avg_loss, acc


def main():
    parser = argparse.ArgumentParser(description='Train ECA-ResNet on CIFAR-100')
    parser.add_argument('--epochs', default=100, type=int)
    parser.add_argument('--batch-size', default=128, type=int)
    parser.add_argument('--lr', default=0.1, type=float)
    parser.add_argument('--resume', action='store_true', help='Resume from checkpoint')
    parser.add_argument('--save-dir', default='./checkpoints', type=str)
    args = parser.parse_args()

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(device)

    print('==> Preparing data..')
    trainloader, testloader = get_cifar100_loaders(batch_size=args.batch_size, num_workers=8)

    print('==> Building model..')
    model = eca_resnet18(num_classes=100)
    model = model.to(device)

    profile_model(model, device, model_name='ECA-ResNet18')

    if args.resume:
        checkpoint_path = os.path.join(args.save_dir, 'latest.pth')
        if os.path.isfile(checkpoint_path):
            print(f"==> Resuming from {checkpoint_path}")
            checkpoint = torch.load(checkpoint_path)
            model.load_state_dict(checkpoint['model_state'])
            start_epoch = checkpoint['epoch']
            best_acc = checkpoint['best_acc']
        else:
            print("No checkpoint found. Starting from scratch.")
            start_epoch = 0
            best_acc = 0.0
    else:
        start_epoch = 0
        best_acc = 0.0

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)

    best_lr = find_lr(model, trainloader, optimizer, criterion, device)

    # 设置为推荐的学习率
    for param_group in optimizer.param_groups:
        param_group['lr'] = best_lr

    os.makedirs(args.save_dir, exist_ok=True)

    train_losses, test_losses = [], []
    train_accs, test_accs = [], []


    for epoch in range(start_epoch, args.epochs):
        print(f"\nEpoch {epoch+1}/{args.epochs}")
        train_loss, train_acc = train(trainloader, model, criterion, optimizer, epoch, device)
        test_loss, test_acc = test(testloader, model, criterion, device)
        scheduler.step()

        print(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
        print(f"Test  Loss: {test_loss:.4f}, Test  Acc: {test_acc:.4f}")

        train_losses.append(train_loss)
        train_accs.append(train_acc)
        test_losses.append(test_loss)
        test_accs.append(test_acc)

        # 保存最佳模型
        is_best = test_acc > best_acc
        if is_best:
            best_acc = test_acc
            torch.save({
                'model_state': model.state_dict(),
                'epoch': epoch,
                'best_acc': best_acc
            }, os.path.join(args.save_dir, 'best.pth'))

        # 保存最后一轮模型
        torch.save({
            'model_state': model.state_dict(),
            'epoch': epoch,
            'best_acc': best_acc
        }, os.path.join(args.save_dir, 'latest.pth'))

    plot_metrics( test_losses, test_accs, save_path=args.save_dir)


if __name__ == '__main__':
    main()
